package main

import (
	"fmt"
	"log"
	"math/rand"
	"time"

	"xagent/internal/actors"
	"xagent/internal/knowledge"
	"xagent/internal/metrics"
	"xagent/internal/optimizer"
	"xagent/internal/scheduler"
	proto "xagent/proto/v1"

	"github.com/asynkron/protoactor-go/actor"
)

func main() {
	// 创建 actor system
	system := actor.NewActorSystem()

	// 创建知识库引擎
	rag := knowledge.NewRAGEngine(&knowledge.Config{
		EmbeddingModel: "text-embedding-ada-002",
		VectorStore:    "inmemory",
	})

	// 创建指标收集器
	collector := metrics.NewCollector()

	// 创建调度器
	resourcePool := scheduler.NewResourcePool(100, 5, 0.8)
	taskScheduler := scheduler.NewScheduler(10, resourcePool)

	// 创建资源优化器
	optimizerConfig := &optimizer.OptimizerConfig{
		MaxLoadPerAgent:      5,
		TargetResponseTime:   time.Second * 2,
		OptimizationPeriod:   time.Second * 10,
		LoadBalanceThreshold: 0.8,
	}
	resourceOptimizer := optimizer.NewResourceOptimizer(taskScheduler, collector, optimizerConfig)

	// 创建工程师团队
	engineers := make([]*actors.AIAgent, 5)
	for i := 0; i < 5; i++ {
		engineers[i] = actors.NewAIAgent(
			fmt.Sprintf("engineer-%d", i+1),
			"Software engineer responsible for implementation")
		engineers[i].SetKnowledgeBase(rag)
	}

	// 创建一批任务
	tasks := []*proto.Task{
		{
			Id:          "task-1",
			Name:        "ImplementUserAuth",
			Description: "实现用户认证系统",
			Status:      proto.TaskStatus_TASK_STATUS_PENDING,
			Tags:        []string{"auth", "security"},
		},
		{
			Id:          "task-2",
			Name:        "ImplementDatabase",
			Description: "实现数据库访问层",
			Status:      proto.TaskStatus_TASK_STATUS_PENDING,
			Tags:        []string{"database", "persistence"},
		},
		{
			Id:          "task-3",
			Name:        "ImplementAPI",
			Description: "实现 REST API",
			Status:      proto.TaskStatus_TASK_STATUS_PENDING,
			Tags:        []string{"api", "rest"},
		},
		{
			Id:          "task-4",
			Name:        "ImplementCache",
			Description: "实现缓存系统",
			Status:      proto.TaskStatus_TASK_STATUS_PENDING,
			Tags:        []string{"cache", "performance"},
		},
		{
			Id:          "task-5",
			Name:        "ImplementLogging",
			Description: "实现日志系统",
			Status:      proto.TaskStatus_TASK_STATUS_PENDING,
			Tags:        []string{"logging", "monitoring"},
		},
	}

	// 启动资源优化器
	resourceOptimizer.Start()

	// 分配初始任务
	for i, task := range tasks {
		engineerIndex := i % len(engineers)
		engineer := engineers[engineerIndex]

		// 更新任务元数据
		task.Metadata = map[string]string{
			"assigned_to": engineer.GetRole(),
			"priority":    "medium",
		}

		// 分配任务
		if err := taskScheduler.MigrateTask(task.Id, "", engineer.GetRole()); err != nil {
			log.Printf("Failed to assign task %s: %v", task.Id, err)
			continue
		}

		// 发送任务给工程师
		msg := &proto.AgentMessage{
			Id:   task.Id,
			Type: proto.MessageType_MESSAGE_TYPE_REQUEST,
			Content: &proto.AgentMessage_Message{
				Message: task.Description,
			},
			Metadata: map[string]string{
				"type":     "task_assignment",
				"task_id":  task.Id,
				"priority": "medium",
			},
		}

		if parent := engineer.GetContext().Parent(); parent != nil {
			engineer.GetContext().Send(parent, msg)
		}
	}

	// 模拟任务执行和负载变化
	go func() {
		for {
			time.Sleep(time.Second * 3)

			// 随机选择一个工程师增加负载
			engineerIndex := rand.Intn(len(engineers))
			engineer := engineers[engineerIndex]

			// 创建新任务
			task := &proto.Task{
				Id:          fmt.Sprintf("dynamic-task-%d", time.Now().UnixNano()),
				Name:        "DynamicTask",
				Description: "动态生成的任务",
				Status:      proto.TaskStatus_TASK_STATUS_PENDING,
				Tags:        []string{"dynamic"},
				Metadata: map[string]string{
					"assigned_to": engineer.GetRole(),
					"priority":    "high",
				},
			}

			// 分配任务
			if err := taskScheduler.MigrateTask(task.Id, "", engineer.GetRole()); err != nil {
				log.Printf("Failed to assign dynamic task: %v", err)
				continue
			}

			// 发送任务给工程师
			msg := &proto.AgentMessage{
				Id:   task.Id,
				Type: proto.MessageType_MESSAGE_TYPE_REQUEST,
				Content: &proto.AgentMessage_Message{
					Message: task.Description,
				},
				Metadata: map[string]string{
					"type":     "task_assignment",
					"task_id":  task.Id,
					"priority": "high",
				},
			}

			if parent := engineer.GetContext().Parent(); parent != nil {
				engineer.GetContext().Send(parent, msg)
			}
		}
	}()

	// 等待一段时间观察负载均衡
	time.Sleep(time.Second * 30)

	// 获取系统负载情况
	systemLoad := taskScheduler.GetSystemLoad()
	log.Printf("Final System Load: %+v", systemLoad)

	// 获取优化指标
	optimizationMetrics := collector.GetMetrics()
	log.Printf("Optimization Metrics: %+v", optimizationMetrics)

	// 停止资源优化器
	resourceOptimizer.Stop()

	// 清理资源
	system.Shutdown()
}
